import json
import os

import pandas as pd
from torchvision import datasets, transforms

if os.name == "nt":
    ROOT_PATH = "D:\projects\datesets"
else:
    ROOT_PATH = "/opt/datesets"


class DataDownloader:
    """
    数据下载器
    """
    pass

    @staticmethod
    def _get_raw_data(data_path: str) -> pd.DataFrame:
        """
        获取原始数据，根据文件类型读取并转换为 DataFrame 格式

        支持的文件类型:
        - .csv
        - .xlsx / .xls
        - .json
        - .jsonl
        - .parquet
        - .txt (默认以制表符分隔)
        """
        if not os.path.exists(data_path):
            raise FileNotFoundError(f"文件不存在: {data_path}")

        ext = os.path.splitext(data_path)[-1].lower()

        if ext == ".csv":
            df = pd.read_csv(data_path)

        elif ext in [".xlsx", ".xls"]:
            df = pd.read_excel(data_path)

        elif ext == ".json":
            # 普通 JSON，可能是 dict 或 list
            with open(data_path, "r", encoding="utf-8") as f:
                data = json.load(f)
            df = pd.json_normalize(data)

        elif ext == ".jsonl":
            df = pd.read_json(data_path, lines=True)

        elif ext == ".parquet":
            df = pd.read_parquet(data_path)

        elif ext == ".txt":
            # 尝试自动识别分隔符（逗号、制表符、分号）
            try:
                df = pd.read_csv(data_path, sep=None, engine="python")
            except Exception:
                df = pd.read_csv(data_path, sep="\t")

        else:
            raise ValueError(f"暂不支持的文件类型: {ext}")

        return df

    @staticmethod
    def get_row_data(data_name):
        """
        原始数据集获取
        """
        if data_name == "mnist":
            train_dataset = datasets.MNIST(
                root=ROOT_PATH,
                train=True,
                download=True,
                transform=None
            )
            test_dataset = datasets.MNIST(
                root=ROOT_PATH,
                train=False,
                download=True,
                transform=None
            )
            return train_dataset, test_dataset
